package org.lanyu.springainovel.common.config;

import lombok.extern.slf4j.Slf4j;
import org.apache.http.Header;
import org.apache.http.HttpHost;
import org.apache.http.message.BasicHeader;
import org.elasticsearch.client.Request;
import org.elasticsearch.client.Response;
import org.elasticsearch.client.RestClient;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.ai.vectorstore.elasticsearch.ElasticsearchVectorStore;
import org.springframework.ai.vectorstore.elasticsearch.ElasticsearchVectorStoreOptions;
import org.springframework.ai.vectorstore.elasticsearch.SimilarityFunction;
import org.springframework.ai.vectorstore.SimpleVectorStore;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.context.annotation.Bean;
import org.springframework.stereotype.Component;

import java.nio.charset.StandardCharsets;
import java.util.Base64;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

/**
 * 向量存储管理�?
 * 用于动态创建和管理Elasticsearch向量存储
 */
@Slf4j
@Component
public class VectorStoreManager {

    @Value("${spring.elasticsearch.uris:http://localhost:9200}")
    private String elasticsearchUris;
    
    @Value("${spring.elasticsearch.username:}")
    private String elasticsearchUsername;
    
    @Value("${spring.elasticsearch.password:}")
    private String elasticsearchPassword;

    @Value("${spring.ai.vectorstore.elasticsearch.index-name:my-index}")
    private String indexName;
    
    @Value("${spring.ai.vectorstore.enabled:true}")
    private boolean vectorStoreEnabled;
    
    private final Map<String, VectorStore> vectorStoreCache = new ConcurrentHashMap<>();
    
    /**
     * 创建手动配置的RestClient
     */
    @Bean
    private RestClient createRestClient() {
        try {
            log.info("创建手动配置的RestClient: URI={}, 用户�?{}", elasticsearchUris, elasticsearchUsername);
            
            // 解析URI
            String uri = elasticsearchUris.split(",")[0]; // 取第一个URI
            HttpHost httpHost = HttpHost.create(uri);
            
            // 创建认证�?
            String credentials = elasticsearchUsername + ":" + elasticsearchPassword;
            String encodedCredentials = Base64.getEncoder().encodeToString(
                credentials.getBytes(StandardCharsets.UTF_8));
            
            Header authHeader = new BasicHeader("Authorization", "Basic " + encodedCredentials);
            
            RestClient restClient = RestClient.builder(httpHost)
                .setDefaultHeaders(new Header[]{authHeader})
                .build();
            Response response = restClient.performRequest(new Request("GET", "/"));
            int statusCode = response.getStatusLine().getStatusCode();
            if (statusCode != 200) {
                log.error("创建RestClient失败: 状态码=" + statusCode);
                throw new RuntimeException("创建RestClient失败: 状态码=" + statusCode);
            }
            log.info("RestClient创建成功");
            return restClient;
        } catch (Exception e) {
            log.error("创建RestClient失败: {}", e.getMessage(), e);
            throw new RuntimeException("创建RestClient失败", e);
        }
    }
    
    /**
     * 确保Elasticsearch索引存在
     */
    private void ensureIndexExists(RestClient restClient, String indexName, int dimensions) {
        try {
            // 检查索引是否存�?
            Request checkRequest = new Request("HEAD", "/" + indexName);
            Response checkResponse = restClient.performRequest(checkRequest);
            
            if (checkResponse.getStatusLine().getStatusCode() == 404) {
                // 索引不存在，创建�?
                log.info("索引 {} 不存在，正在创建...", indexName);
                
                Request createRequest = new Request("PUT", "/" + indexName);
                
                // 创建索引的mapping配置，包含向量字�?
                String mappingJson = "{"
                    + "\"mappings\":{"
                    + "\"properties\":{"
                    + "\"embedding\":{"
                    + "\"type\":\"dense_vector\","
                    + "\"dims\":" + dimensions + ","
                    + "\"similarity\":\"cosine\","
                    + "\"index\":true"
                    + "},"
                    + "\"content\":{"
                    + "\"type\":\"text\","
                    + "\"analyzer\":\"standard\""
                    + "},"
                    + "\"metadata\":{"
                    + "\"type\":\"object\","
                    + "\"enabled\":true"
                    + "}"
                    + "}"
                    + "}"
                    + "}";
                
                createRequest.setJsonEntity(mappingJson);
                Response createResponse = restClient.performRequest(createRequest);
                
                if (createResponse.getStatusLine().getStatusCode() == 200) {
                    log.info("索引 {} 创建成功", indexName);
                } else {
                    log.error("索引 {} 创建失败: 状态码={}", indexName, createResponse.getStatusLine().getStatusCode());
                }
            } else {
                log.info("索引 {} 已存在", indexName);
            }
        } catch (Exception e) {
            log.error("检查或创建索引失败: {} - {}", indexName, e.getMessage(), e);
            throw new RuntimeException("检查或创建索引失败: " + indexName, e);
        }
    }
    
    /**
     * 创建Elasticsearch向量存储
     */
    private VectorStore createElasticsearchVectorStore(EmbeddingModel embeddingModel, String indexName, int dimensions) {
        try {
            log.info("开始创建Elasticsearch向量存储: 索引={}, 维度={}", indexName, dimensions);
            
            // 创建选项
            ElasticsearchVectorStoreOptions options = new ElasticsearchVectorStoreOptions();
            options.setIndexName(indexName);
            options.setDimensions(dimensions);
            options.setSimilarity(SimilarityFunction.cosine);
            
            // 使用手动配置的RestClient
            RestClient restClient = createRestClient();
            
            // 首先确保索引存在
            ensureIndexExists(restClient, indexName, dimensions);
            
            VectorStore vectorStore = ElasticsearchVectorStore.builder(restClient, embeddingModel)
                .options(options)
                .initializeSchema(false)  // 修改：禁用自动初始化，因为我们已经手动创建了索引
                .build();
            
            // 检查索引是否已存在，如果存在则跳过验证
            boolean indexExists = checkIndexExists(restClient, indexName);
            
            if (!indexExists) {
                // 索引不存在，需要初始化
                try {
                    vectorStore.add(java.util.Collections.emptyList());
                    log.info("Elasticsearch向量存储索引创建验证成功: {} - 维度: {}", indexName, dimensions);
                } catch (Exception e) {
                    log.warn("首次索引创建验证失败，尝试强制初始化: {} - {}", indexName, e.getMessage());
                    // 如果失败，尝试强制初始化 - 注意这里需要获取具体的ElasticsearchVectorStore实例
                    if (vectorStore instanceof ElasticsearchVectorStore) {
                        try {
                            ElasticsearchVectorStore esVectorStore =
                                (ElasticsearchVectorStore) vectorStore;
                            java.lang.reflect.Method initMethod = esVectorStore.getClass().getDeclaredMethod("initializeSchema");
                            initMethod.setAccessible(true);
                            initMethod.invoke(esVectorStore);
                            log.info("强制初始化索引成功: {}", indexName);
                        } catch (Exception initException) {
                            log.error("强制初始化索引失败: {} - {}", indexName, initException.getMessage());
                        }
                    } else {
                        log.error("向量存储不是ElasticsearchVectorStore类型，无法强制初始化: {}", vectorStore.getClass().getName());
                    }
                }
            } else {
                log.info("索引 {} 已存在，跳过初始化验证", indexName);
            }
            
            log.info("创建Elasticsearch向量存储成功: {} - 维度: {}", indexName, dimensions);
            return vectorStore;
        } catch (Exception e) {
            log.error("创建Elasticsearch向量存储失败: {} - {}", indexName, e.getMessage(), e);
            throw new RuntimeException("创建Elasticsearch向量存储失败", e);
        }
    }
    
    /**
     * 检查索引是否存在
     */
    private boolean checkIndexExists(RestClient restClient, String indexName) {
        try {
            Request checkRequest = new Request("HEAD", "/" + indexName);
            Response checkResponse = restClient.performRequest(checkRequest);
            return checkResponse.getStatusLine().getStatusCode() == 200;
        } catch (Exception e) {
            log.debug("检查索引是否存在时出错: {} - {}", indexName, e.getMessage());
            return false;
        }
    }
    
    /**
     * 获取或创建向量存储（带缓存）
     * 
     * @param embeddingModel 嵌入模型
     * @param indexName 索引名称
     * @param dimensions 向量维度
     * @return 向量存储实例
     */
    public VectorStore getOrCreateVectorStore(EmbeddingModel embeddingModel, 
                                             String indexName, 
                                             int dimensions) {
        // 如果向量存储功能被禁用，返回一个简单的内存向量存储
        if (!vectorStoreEnabled) {
            log.info("ES向量存储功能已禁用，使用内存向量存储");
            return SimpleVectorStore.builder(embeddingModel).build();
        }
        
        String cacheKey = indexName + "_" + dimensions;
        return vectorStoreCache.computeIfAbsent(cacheKey, 
            k -> createElasticsearchVectorStore(embeddingModel, indexName, dimensions));
    }
    
    /**
     * 获取默认向量存储
     * 
     * @param embeddingModel 嵌入模型
     * @return 向量存储实例
     */
    public VectorStore getDefaultVectorStore(EmbeddingModel embeddingModel) {
        // 如果向量存储功能被禁用，返回一个简单的内存向量存储
        if (!vectorStoreEnabled) {
            log.info("ES向量存储功能已禁用，使用内存向量存储");
            return SimpleVectorStore.builder(embeddingModel).build();
        }
        
        return getOrCreateVectorStore(embeddingModel, indexName, 1024);
    }
    
    /**
     * 清除指定索引的缓�?
     * 
     * @param indexName 索引名称
     * @param dimensions 向量维度
     */
    public void clearCache(String indexName, int dimensions) {
        String cacheKey = indexName + "_" + dimensions;
        vectorStoreCache.remove(cacheKey);
        log.info("清除向量存储缓存: {}", cacheKey);
    }
    
    /**
     * 清除所有缓存
     */
    public void clearAllCache() {
        vectorStoreCache.clear();
        log.info("清除所有向量存储缓存");
    }
    
    /**
     * 检查向量存储是否启用
     * 
     * @return true表示启用，false表示禁用
     */
    public boolean isVectorStoreEnabled() {
        return vectorStoreEnabled;
    }
}
